#!/usr/bin/env python3
"""
Gripper loading and point cloud generation utilities.

This module provides functionality to:
1. Load embodiment configurations based on character codes
2. Parse URDF files to extract gripper joint and link information
3. Generate point clouds on gripper surfaces with finger opening/closing control
4. Apply proper transformations based on gripper joint configurations

The system supports the following embodiments (excluding curobo99 and ur5-franka-gripper):
- aloha-agilex (a)
- piper (p) 
- ur5-wsg (u)
- franka-panda (f)
- ARX-X5 (x)
"""
import numpy as np
import json
import os
import yaml
import warnings
import xml.etree.ElementTree as ET
from typing import Dict, List, Tuple, Optional, Any
from scipy.spatial.transform import Rotation
import glob
import trimesh
import trimesh.util
from trimesh.sample import sample_surface
import inspect
from pathlib import Path
from scalePcd import scalePcd

warnings.filterwarnings('ignore')

def get_all_attributes(obj, max_depth=2, current_depth=0):
    """
    Get all attributes of an unknown object recursively.
    
    Args:
        obj: The object to inspect
        max_depth: Maximum recursion depth to avoid infinite loops
        current_depth: Current recursion depth
        
    Returns:
        dict: Dictionary containing all attributes and their values/types
    """
    if obj is None:
        return {"type": "None", "value": None}
    
    if current_depth >= max_depth:
        return {"type": str(type(obj)), "value": "Max depth reached"}
    
    result = {
        "type": str(type(obj)),
        "attributes": {},
        "methods": {},
        "properties": {}
    }
    
    try:
        # Get all attributes
        for attr_name in dir(obj):
            if attr_name.startswith('_'):  # Skip private attributes
                continue
                
            try:
                attr_value = getattr(obj, attr_name)
                
                # Skip methods and functions
                if callable(attr_value):
                    result["methods"][attr_name] = {
                        "type": str(type(attr_value)),
                        "doc": getattr(attr_value, '__doc__', 'No docstring')
                    }
                    continue
                
                # Handle different types of attributes
                if isinstance(attr_value, (int, float, str, bool, type(None))):
                    result["attributes"][attr_name] = {
                        "type": str(type(attr_value)),
                        "value": attr_value
                    }
                elif isinstance(attr_value, (list, tuple)):
                    result["attributes"][attr_name] = {
                        "type": str(type(attr_value)),
                        "length": len(attr_value),
                        "sample": str(attr_value[:3]) if len(attr_value) > 0 else "empty"
                    }
                elif isinstance(attr_value, dict):
                    result["attributes"][attr_name] = {
                        "type": str(type(attr_value)),
                        "keys": list(attr_value.keys())[:5],  # Show first 5 keys
                        "length": len(attr_value)
                    }
                elif isinstance(attr_value, np.ndarray):
                    result["attributes"][attr_name] = {
                        "type": "numpy.ndarray",
                        "shape": attr_value.shape,
                        "dtype": str(attr_value.dtype),
                        "sample": str(attr_value.flatten()[:5]) if attr_value.size > 0 else "empty"
                    }
                else:
                    # For complex objects, try to get basic info
                    try:
                        result["attributes"][attr_name] = {
                            "type": str(type(attr_value)),
                            "str_repr": str(attr_value)[:100] + "..." if len(str(attr_value)) > 100 else str(attr_value)
                        }
                    except:
                        result["attributes"][attr_name] = {
                            "type": str(type(attr_value)),
                            "value": "Could not get string representation"
                        }
                        
            except Exception as e:
                result["attributes"][attr_name] = {
                    "type": "Error",
                    "error": str(e)
                }
                
    except Exception as e:
        result["error"] = f"Error inspecting object: {str(e)}"
    
    return result

def print_object_info(obj, name="object"):
    """
    Print comprehensive information about an object in a readable format.
    
    Args:
        obj: The object to inspect
        name: Name of the object for display purposes
    """
    print(f"\n=== {name} Information ===")
    info = get_all_attributes(obj)
    
    print(f"Type: {info['type']}")
    
    if "error" in info:
        print(f"Error: {info['error']}")
        return
    
    if info["attributes"]:
        print(f"\nAttributes ({len(info['attributes'])}):")
        for attr_name, attr_info in info["attributes"].items():
            print(f"  {attr_name}: {attr_info}")
    
    if info["methods"]:
        print(f"\nMethods ({len(info['methods'])}):")
        for method_name, method_info in info["methods"].items():
            print(f"  {method_name}: {method_info['type']}")
            if method_info['doc']:
                doc_lines = method_info['doc'].strip().split('\n')
                for line in doc_lines[:2]:  # Show first 2 lines of docstring
                    print(f"    {line}")
    
    if info["properties"]:
        print(f"\nProperties ({len(info['properties'])}):")
        for prop_name, prop_info in info["properties"].items():
            print(f"  {prop_name}: {prop_info}")

# Global cache for loaded gripper data
_GRIPPER_CACHE = {}

# Paths
EMBODIMENT_MAPPING_PATH = str(Path(__file__).parent / 'embodiment_mapping.json')
EMBODIMENT_CONFIG_PATH = str(Path(__file__).parent.parent / 'RoboTwin/task_config/_embodiment_config.yml')
MAPPING = json.load(open(EMBODIMENT_MAPPING_PATH, 'r'))
CONFIG = yaml.safe_load(open(EMBODIMENT_CONFIG_PATH, 'r'))

SHOW_HAND = True
SHOW_L_FINGER = False
SHOW_R_FINGER = False

class GripperLoader:
    def __init__(self):
        self.full_urdf_path = None
        self.config = None
        self.use_real_gripper = True
        self.gripper_mesh = None
        self.gripper_mesh_path = str(Path(__file__).parent.parent / 'RoboTwin/assets/embodiments/franka-panda/franka_description/meshes/visual/hand.glb')
        
    def load_gripper_mesh(self):
        """Load the gripper mesh from GLB file and cache it."""
        if self.gripper_mesh is None:
            assert os.path.exists(self.gripper_mesh_path), f"Gripper mesh file not found: {self.gripper_mesh_path}"
            loaded_mesh = trimesh.load(self.gripper_mesh_path)
            
            # Handle case where GLB file loads as a Scene (containing multiple meshes)
            if isinstance(loaded_mesh, trimesh.Scene):
                # Get all meshes from the scene
                meshes = list(loaded_mesh.geometry.values())
                if len(meshes) == 0:
                    raise ValueError(f"No meshes found in scene: {self.gripper_mesh_path}")
                elif len(meshes) == 1:
                    self.gripper_mesh = meshes[0]
                else:
                    # If multiple meshes, combine them into a single mesh
                    self.gripper_mesh = trimesh.util.concatenate(meshes)
            else:
                # If it's already a single mesh
                self.gripper_mesh = loaded_mesh
                
            print(f"Successfully loaded gripper mesh from {self.gripper_mesh_path}")
        return True
    def reinit(self, embodiment_code: str):
        if self.full_urdf_path is None or self.config is None:
            self.full_urdf_path = self.get_urdf_file_from_code(embodiment_code)
            self.config = self.get_embodiment_config(embodiment_code)
    def get_urdf_file_from_code(self,embodiment_code: str) -> Optional[str]:
        """
        Given an embodiment code, return the path to the main URDF file for that embodiment.
        Args:
            embodiment_code: Single character code (a, p, u, f, x)
        Returns:
            Path to the main URDF file or None if not found
        """
        if embodiment_code not in MAPPING:
            assert False, f"Unknown embodiment code '{embodiment_code}'"
        embodiment_name = MAPPING[embodiment_code]
        if embodiment_name not in CONFIG:
            assert False, f"No configuration found for embodiment '{embodiment_name}'"
        embodiment_path = Path(__file__).parent.parent / "RoboTwin" / CONFIG[embodiment_name].get('file_path')
        if not embodiment_path or not os.path.exists(embodiment_path):
            assert False, f"Embodiment path does not exist: {embodiment_path}"
        
        # Get the config file to find the URDF path
        config_file = os.path.join(embodiment_path, 'config.yml')
        if not os.path.exists(config_file):
            assert False, f"Config file not found: {config_file}"
        
        config = yaml.safe_load(open(config_file, 'r'))
        urdf_path = config.get('urdf_path', '*.urdf')
        
        # Resolve relative path
        if urdf_path.startswith('./'):
            urdf_path = urdf_path[2:]
        
        full_urdf_path = os.path.join(embodiment_path, urdf_path)
        if not os.path.exists(full_urdf_path):
            # Try to find URDF file
            urdf_candidates = [
                '*.urdf',
                'urdf/*.urdf',
                'franka_description/*.urdf'
            ]
            for pattern in urdf_candidates:
                matches = glob.glob(os.path.join(embodiment_path, pattern))
                if matches:
                    non_tmp_files = [f for f in matches if 'tmp' not in f.lower()]
                    if non_tmp_files:
                        return non_tmp_files[0]
                    return matches[0]
            assert False, f"Could not find URDF file in {embodiment_path}"
        return full_urdf_path
    def get_embodiment_config(self,embodiment_code: str) -> Dict[str, Any]:
        """
        Get the full configuration for an embodiment.
        Args:
            embodiment_code: Single character embodiment code
        Returns:
            Dictionary containing embodiment configuration
        """
        if embodiment_code not in MAPPING:
            assert False, f"Unknown embodiment code '{embodiment_code}'"
        embodiment_name = MAPPING[embodiment_code]
        embodiment_path = Path(__file__).parent.parent / "RoboTwin" / CONFIG[embodiment_name].get('file_path')
        config_file = os.path.join(embodiment_path, 'config.yml')
        if not os.path.exists(config_file):
            assert False, f"Config file not found: {config_file}"
        return yaml.safe_load(open(config_file, 'r'))
    def quat_to_rotmat(self,q):
        w, x, y, z = q
        return np.array([
            [1-2*y*y-2*z*z, 2*x*y-2*z*w, 2*x*z+2*y*w],
            [2*x*y+2*z*w, 1-2*x*x-2*z*z, 2*y*z-2*x*w],
            [2*x*z-2*y*w, 2*y*z+2*x*w, 1-2*x*x-2*y*y]
        ])
    def get_rectangle_pcd(self, pose, num_points, left=True):
        box_size = np.array([0.08, 0.02, 0.04])
        box_center = np.array([0, 0, 0])
        points_local = np.random.uniform(-0.5, 0.5, (num_points, 3)) * box_size + box_center
        rot = self.quat_to_rotmat(pose[3:7])
        
        # Apply delta_matrix transformation to correct gripper orientation
        if self.config is not None and 'delta_matrix' in self.config:
            delta_matrix = np.array(self.config['delta_matrix'])
            # Apply delta_matrix to the rotation matrix
            rot = rot @ delta_matrix
        else:
            print(f"get_rectangle_pcd No delta_matrix found in {self.config}")
        
        points_world = points_local @ rot.T + pose[0:3]
        color = np.array([128, 0, 0], dtype=np.float32) if left else np.array([0, 128, 0], dtype=np.float32)
        # Colors are already in 0-255 range, no normalization needed
        print(f"loadGripper.py: get_rectangle_pcd - Generated color: {color} (0-255 range)")
        colors = np.tile(color, (num_points, 1))
        return np.concatenate([points_world, colors], axis=1)
    def get_real_gripper_pcd(self,n: int, pose: np.ndarray, finger_dist: float = 0.5) -> np.ndarray:
        """
        Generate point cloud from real gripper mesh by sampling points on its surface.
        
        Args:
            n: Number of points to sample
            pose: 7D pose array [x, y, z, qw, qx, qy, qz]
            finger_dist: Finger distance (not used for mesh sampling but kept for compatibility)
            
        Returns:
            np.ndarray: Point cloud with colors, shape (n, 6) [x, y, z, r, g, b]
        """
        # Load the gripper mesh if not already loaded
        self.load_gripper_mesh()
        points, face_indices = sample_surface(self.gripper_mesh, n)
        
        # Get colors from the mesh
        if not (self.gripper_mesh is not None and 
            hasattr(self.gripper_mesh, 'visual') and 
            self.gripper_mesh.visual is not None):
            print_object_info(self.gripper_mesh, "gripper_mesh")
            exit()
            
        visual = self.gripper_mesh.visual
        
        # If it's TextureVisuals, convert to ColorVisuals first
        if hasattr(visual, 'kind') and visual.kind == 'texture':
            color_visual = visual.to_color()
            
            if not hasattr(color_visual, 'vertex_colors'):
                print_object_info(color_visual, "color_visual")
                exit()
            
            vertex_colors = color_visual.vertex_colors
            face_vertices = self.gripper_mesh.faces[face_indices]
            colors = np.mean(vertex_colors[face_vertices], axis=1)[:, :3].astype(np.float32)
        
        # If it's already ColorVisuals, access vertex_colors directly
        elif hasattr(visual, 'vertex_colors'):
            vertex_colors = visual.vertex_colors
            face_vertices = self.gripper_mesh.faces[face_indices]
            colors = np.mean(vertex_colors[face_vertices], axis=1)[:, :3].astype(np.float32)
        else:
            print_object_info(visual, "visual")
            exit()
        
        # Apply pose transformation with delta_matrix correction
        rot = self.quat_to_rotmat(pose[3:7])
        
        # Apply delta_matrix transformation to correct gripper orientation
        if self.config is not None and 'delta_matrix' in self.config:
            delta_matrix = np.array(self.config['delta_matrix'])
            # Apply delta_matrix to the rotation matrix
            rot = rot @ delta_matrix
        else:
            print(f"get_real_gripper_pcd No delta_matrix found in {self.config}")
        
        points_world = points @ rot.T + pose[0:3]
        
        # Combine points and colors
        pcd = np.concatenate([points_world, colors], axis=1)
        pcd = scalePcd(pcd, tag="get_real_gripper_pcd")
        return pcd
        
    def get_two_gripper_pcd(self,gripper_key: str, n1: int, n2: int, 
                    pose1: np.ndarray, pose2: np.ndarray, 
                    finger_dist: float = 0.5) -> np.ndarray:
        if self.use_real_gripper:
            pcd1 = self.get_real_gripper_pcd(n1,pose1,finger_dist)
            pcd2 = self.get_real_gripper_pcd(n2,pose2,finger_dist)
        else:
            pcd1 = self.get_rectangle_pcd(pose1, n1, left=True)
            pcd2 = self.get_rectangle_pcd(pose2, n2, left=False)
        return np.concatenate([pcd1, pcd2], axis=0)

gripper_loader = GripperLoader()
def get_two_gripper_pcd(gripper_key, n, pose1, pose2, finger_dist=0.5):
    """
    Global function to get point clouds for both grippers.
    
    Args:
        gripper_key: Gripper identifier
        n: Total number of points for both grippers (excludes finger points)
        pose1: 7D pose for left gripper [x, y, z, qw, qx, qy, qz]
        pose2: 7D pose for right gripper [x, y, z, qw, qx, qy, qz]
        finger_dist: Finger distance (optional, for compatibility)
    
    Returns:
        np.ndarray: Combined point cloud for both grippers
    """
    gripper_loader.reinit(gripper_key)
    
    # Split points equally between left and right grippers
    n1 = n // 2
    n2 = n - n1
    
    return gripper_loader.get_two_gripper_pcd(gripper_key, n1, n2, pose1, pose2, finger_dist)
# if __name__ == "__main__":
#     gripper_loader.reinit('f')
#     pcd = gripper_loader.get_real_gripper_pcd(1000, np.array([0, 0, 0, 1, 0, 0, 0]))
#     print(pcd.shape)
